# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: taekwon

import sys
import os
import json
import pickle
import logging
import collections
import datetime
import traceback
import pandas

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor
from coin.base.query_util import query_pta
import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.tool.pnl_plot as pnl_plot_job


def main(_):
  current_time = datetime.datetime.utcnow()
  yesterday = current_time.date() - datetime.timedelta(days=1)
  target_exchanges = flags.FLAGS.exchanges.split(',')

  stratdf = query_pta(yesterday,
                      yesterday + datetime.timedelta(days=1),
                      business_units=['Coin'],
                      exchanges=target_exchanges)
  stratdf = stratdf[
      ['strategy_name', 'account_id', 'owner', 'trader', 'exchange', 'market_type', 'machine']].drop_duplicates()

  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()
  trading_dates = abase.get_trading_dates(baseflags.trading_date)

  for _, rows in stratdf.iterrows():
    osflags.orderlog_machine = rows.machine
    osflags.strategy_name = rows.strategy_name
    args = tuple([baseflags, fsflags, osflags])
    rettuples = []
  
    if flags.FLAGS.cpu == 1 or len(trading_dates) == 1:
      for trading_date in trading_dates:
        rettuples.append(pnl_plot_job.launch(trading_date, *args))
    else:
      with ProcessPoolExecutor(max_workers=flags.FLAGS.cpu) as executor:
        futures = []
        for trading_date in trading_dates:
          futures.append(executor.submit(pnl_plot_job.launch, trading_date, *args))
        for future in futures:
          res = future.result()
          if res is not None:
            rettuples.append(res)

    orddfs = []
    for orddf, _ in rettuples:
      if orddf is None:
        continue
      orddfs.append(orddf)

    orddf = pandas.concat(orddfs, axis=0, sort=False).reset_index(drop=True)
    orddf = orddf.sort_values('timestamp').reset_index(drop=True)
    resdf = orddf[(orddf.fill_qty > 0.0) & (orddf.fill_price > 0.0)]
    resdf = resdf[['timestamp', 'market_type', 'exchange_type', 'symbol', 'fill_price', 'fill_qty', 'fill_type']]
    resdf['timestamp'] = pandas.DatetimeIndex(resdf['timestamp']).tz_localize(datetime.timezone.utc)
    resdf['timestamp'] = resdf['timestamp'].dt.tz_convert('Asia/Seoul')
    resdf.to_csv(f'{rows.strategy_name}_{baseflags.trading_date}.csv', index=False)


def init_flags():
  pnl_plot_job.init_flags()
  flags.FLAGS.strategy_name = ''
  flags.FLAGS.feed_machine = 'feed-05.ap-northeast-1.aws'
  flags.DEFINE_string('exchanges', 'Korbit,Upbit,Bithumb', 'target exchanges')


if __name__ == '__main__':
  try:
    init_flags()
    app.run(main)
  except Exception:
    traceback.print_exc()
    print(' '.join(sys.argv))

"""
./pyrunner python/coin/strategy/mm/tool/fill_ledger.py --trading_date 20220424-20220430 --exchange=Korbit,Upbit,Bithumb --cpu=3
"""
